import os
import random

def get_file(dir_path):
    file_list = []
    r = ''
    for root, dirs, files in os.walk(dir_path):
        for file in files:
            if file.endswith("jpg") or file.endswith("png") or file.endswith("txt"):
                file_list.append(file.split('.')[0])
    random.shuffle(file_list)
    return file_list

# 将original_path的东西移动到destination_path下
import os
import shutil


def mycopyfile(srcfile, dstpath):  # 复制函数
    if not os.path.isfile(srcfile):
        print("%s not exist!" % (srcfile))
    else:
        fpath, fname = os.path.split(srcfile)  # 分离文件名和路径
        if not os.path.exists(dstpath):
            os.makedirs(dstpath)  # 创建路径
        shutil.copy(srcfile, dstpath)  # 复制文件
        #print("copy %s -> %s" % (srcfile, dstpath))


if __name__ == '__main__':
    train_p = 0.6
    val_p = 0.2
    test_p = 0.2
    path_labels = '''/datasets/uniuav/set4\\labels'''
    path = '''/datasets/uniuav/set4\\images'''
    file_list = get_file(path)
    total_len = len(file_list)+1
    train_list = file_list[0: int(total_len*train_p)]
    val_list = file_list[int(total_len*train_p):int( total_len*(train_p+val_p))]
    test_list = file_list[int( total_len*(train_p+val_p)):]

    train_path = path + "train"
    val_path = path + "val"
    test_path = path + "test"

    for file in train_list:
        mycopyfile(os.path.join(path,file+".jpg"),os.path.join(path,"train"))
        mycopyfile(os.path.join(path_labels, file + ".txt"), os.path.join(path_labels, "train"))
    for file in val_list:
        mycopyfile(os.path.join(path, file + ".jpg"), os.path.join(path, "val"))
        mycopyfile(os.path.join(path_labels, file + ".txt"), os.path.join(path_labels, "val"))
    for file in test_list:
        mycopyfile(os.path.join(path, file + ".jpg"), os.path.join(path, "test"))
        mycopyfile(os.path.join(path_labels, file + ".txt"), os.path.join(path_labels, "test"))




